from openai import OpenAI
import json

# 给AI发送消息
def send_messages(messages):
    # 获得回答
    # 如果模型需要调用工具，则返回的消息会包含 tool_calls 属性
    response = client.chat.completions.create(
        model="deepseek-chat",
        messages=messages,
        tools=tools # 工具列表
    )
    return response.choices[0].message

client = OpenAI(
    api_key='sk-6623ff8a31894caa8a3dc2a2802a94cf',
    base_url="https://api.deepseek.com",
)

def get_weather(city: str) -> str:
    weather_data = {
        "beijing": {
            "location": "Beijing",
            "temperature": {
                "current": 32,
                "low": 26,
                "high": 35
            },
            "rain_probability": 10,   # 百分比
            "humidity": 40  # 百分比
        },
        "shenzhen": {
            "location": "Shenzhen",
            "temperature": {
                "current": 28,
                "low": 24,
                "high": 31
            },
            "rain_probability": 90,   # 百分比
            "humidity": 85     # 百分比
        }
    }

    city_key = city.lower()
    if city_key in weather_data:
        return json.dumps(weather_data[city_key], ensure_ascii=False)
    return json.dumps({"error": "Weather Unavailable"}, ensure_ascii=False)

# 定义工具映射
tool_functions = {
    "get_weather": get_weather,
    # 可以继续添加其他工具函数
}

# 工具列表
tools = [
    {
        "type": "function",
        "function": {
            "name": "get_weather",
            "description": "Get weather of an location, the user shoud supply a location first",
            "parameters": {
                "type": "object",
                "properties": {
                    "location": {
                        "type": "string",
                        "description": "The city, e.g. San Francisco",
                    }
                },
                "required": ["location"]
            },
        }
    },
]

messages = [{"role": "user", "content": "查找shenzhen的天气，然后用一句话告诉我出门要不要带伞"}]

# 将 messages 传入大模型，大模型给出回答
message = send_messages(messages)
# print(message)
print(f"User>\t {messages[0]['content']}")

# 解析模型建议调用的工具
tool = message.tool_calls[0]
func_name = tool.function.name # 函数名称
args = json.loads(tool.function.arguments)
location = args["location"]
print(location)

# 将大模型的回答添加到 messages
messages.append(message)

'''
此时 messages 被添加了自己第一次的回答，可以输出看看
大模型只会告诉你需要调用哪个工具，但不会实际执行工具
'''

print(message)

# 通过名称调用工具函数
result = tool_functions[func_name](location)

messages.append({
    "role": "tool", 
    "tool_call_id": tool.id, 
    "content": result
})

message = send_messages(messages)
print(f"Model>\t {message.content}")